#!/usr/bin/env python3
"""
ES 与 MySQL 数据一致性对账工具

功能：
1. 比对 MySQL user 表与 Elasticsearch user 索引的数据差异
2. 输出不一致的记录（缺失、多余、字段不匹配）
3. 可选：自动修复不一致数据
"""

import argparse
import sys
import json
from typing import Dict, List, Tuple, Set

try:
    import mysql.connector
    from elasticsearch import Elasticsearch
except ImportError:
    print("错误：缺少依赖库，请安装：")
    print("  pip3 install mysql-connector-python elasticsearch")
    sys.exit(1)


class ConsistencyChecker:
    def __init__(self, mysql_config: Dict, es_config: Dict):
        self.mysql_conn = mysql.connector.connect(**mysql_config)
        self.mysql_cursor = self.mysql_conn.cursor(dictionary=True)
        self.es_client = Elasticsearch([es_config['host']])
        self.es_index = es_config.get('index', 'user')

    def fetch_mysql_users(self) -> Dict[str, Dict]:
        """从 MySQL 获取所有用户"""
        query = """
            SELECT user_id, nickname, phone, description, avatar_id
            FROM user
            WHERE user_id IS NOT NULL
        """
        self.mysql_cursor.execute(query)
        users = {}
        for row in self.mysql_cursor.fetchall():
            user_id = row['user_id']
            users[user_id] = {
                'nickname': row['nickname'] or '',
                'phone': row['phone'] or '',
                'description': row['description'] or '',
                'avatar_id': row['avatar_id'] or ''
            }
        return users

    def fetch_es_users(self) -> Dict[str, Dict]:
        """从 ES 获取所有用户"""
        users = {}
        try:
            response = self.es_client.search(
                index=self.es_index,
                body={
                    "query": {"match_all": {}},
                    "size": 10000
                }
            )
            for hit in response['hits']['hits']:
                if hit['_id'] == 'default_index_id':
                    continue
                source = hit['_source']
                user_id = source.get('user_id', '')
                if user_id:
                    users[user_id] = {
                        'nickname': source.get('nickname', ''),
                        'phone': source.get('phone', ''),
                        'description': source.get('description', ''),
                        'avatar_id': source.get('avatar_id', '')
                    }
        except Exception as e:
            print(f"警告：读取 ES 索引失败: {e}")
        return users

    def compare(self) -> Tuple[Set[str], Set[str], Dict[str, List[str]]]:
        """
        比对数据差异
        返回: (仅在MySQL, 仅在ES, 字段不一致的用户)
        """
        mysql_users = self.fetch_mysql_users()
        es_users = self.fetch_es_users()

        mysql_ids = set(mysql_users.keys())
        es_ids = set(es_users.keys())

        only_in_mysql = mysql_ids - es_ids
        only_in_es = es_ids - mysql_ids

        field_mismatches = {}
        for user_id in mysql_ids & es_ids:
            mysql_data = mysql_users[user_id]
            es_data = es_users[user_id]
            diff_fields = []
            for field in ['nickname', 'phone', 'description', 'avatar_id']:
                if mysql_data.get(field, '') != es_data.get(field, ''):
                    diff_fields.append(f"{field}: MySQL='{mysql_data.get(field)}' vs ES='{es_data.get(field)}'")
            if diff_fields:
                field_mismatches[user_id] = diff_fields

        return only_in_mysql, only_in_es, field_mismatches

    def fix_missing_in_es(self, user_ids: Set[str], mysql_users: Dict[str, Dict]) -> int:
        """修复 ES 中缺失的用户"""
        fixed = 0
        for user_id in user_ids:
            data = mysql_users[user_id]
            try:
                self.es_client.index(
                    index=self.es_index,
                    id=user_id,
                    body={
                        'user_id': user_id,
                        'nickname': data['nickname'],
                        'phone': data['phone'],
                        'description': data['description'],
                        'avatar_id': data['avatar_id']
                    }
                )
                fixed += 1
                print(f"  ✓ 修复: 新增用户 {user_id} 到 ES")
            except Exception as e:
                print(f"  ✗ 失败: 新增用户 {user_id} 失败 - {e}")
        return fixed

    def fix_extra_in_es(self, user_ids: Set[str]) -> int:
        """删除 ES 中多余的用户"""
        fixed = 0
        for user_id in user_ids:
            try:
                self.es_client.delete(index=self.es_index, id=user_id)
                fixed += 1
                print(f"  ✓ 修复: 删除 ES 中多余用户 {user_id}")
            except Exception as e:
                print(f"  ✗ 失败: 删除用户 {user_id} 失败 - {e}")
        return fixed

    def fix_field_mismatches(self, mismatches: Dict[str, List[str]], mysql_users: Dict[str, Dict]) -> int:
        """修复字段不一致"""
        fixed = 0
        for user_id, diff_fields in mismatches.items():
            data = mysql_users[user_id]
            try:
                self.es_client.index(
                    index=self.es_index,
                    id=user_id,
                    body={
                        'user_id': user_id,
                        'nickname': data['nickname'],
                        'phone': data['phone'],
                        'description': data['description'],
                        'avatar_id': data['avatar_id']
                    }
                )
                fixed += 1
                print(f"  ✓ 修复: 更新用户 {user_id} 的字段")
            except Exception as e:
                print(f"  ✗ 失败: 更新用户 {user_id} 失败 - {e}")
        return fixed

    def close(self):
        if self.mysql_cursor:
            self.mysql_cursor.close()
        if self.mysql_conn:
            self.mysql_conn.close()


def main():
    parser = argparse.ArgumentParser(description='ES 与 MySQL 数据一致性对账工具')
    parser.add_argument('--mysql-host', default='127.0.0.1', help='MySQL 主机地址')
    parser.add_argument('--mysql-port', type=int, default=3306, help='MySQL 端口')
    parser.add_argument('--mysql-user', default='root', help='MySQL 用户名')
    parser.add_argument('--mysql-password', default='ok', help='MySQL 密码')
    parser.add_argument('--mysql-db', default='suye', help='MySQL 数据库名')
    parser.add_argument('--es-host', default='http://localhost:9200', help='ES 主机地址')
    parser.add_argument('--es-index', default='user', help='ES 索引名称')
    parser.add_argument('--fix', action='store_true', help='自动修复不一致数据')
    parser.add_argument('--fix-missing', action='store_true', help='仅修复 ES 缺失数据')
    parser.add_argument('--fix-extra', action='store_true', help='仅修复 ES 多余数据')
    parser.add_argument('--fix-fields', action='store_true', help='仅修复字段不一致')

    args = parser.parse_args()

    mysql_config = {
        'host': args.mysql_host,
        'port': args.mysql_port,
        'user': args.mysql_user,
        'password': args.mysql_password,
        'database': args.mysql_db
    }

    es_config = {
        'host': args.es_host,
        'index': args.es_index
    }

    checker = ConsistencyChecker(mysql_config, es_config)

    print("=" * 60)
    print("ES 与 MySQL 数据一致性对账")
    print("=" * 60)

    only_mysql, only_es, field_diff = checker.compare()

    print(f"\n📊 对账结果：")
    print(f"  - MySQL 用户总数: {len(checker.fetch_mysql_users())}")
    print(f"  - ES 用户总数: {len(checker.fetch_es_users())}")
    print(f"  - 仅在 MySQL: {len(only_mysql)}")
    print(f"  - 仅在 ES: {len(only_es)}")
    print(f"  - 字段不一致: {len(field_diff)}")

    if only_mysql:
        print(f"\n⚠️  ES 中缺失的用户 ({len(only_mysql)} 个):")
        for user_id in list(only_mysql)[:10]:
            print(f"  - {user_id}")
        if len(only_mysql) > 10:
            print(f"  ... 还有 {len(only_mysql) - 10} 个")

    if only_es:
        print(f"\n⚠️  ES 中多余的用户 ({len(only_es)} 个):")
        for user_id in list(only_es)[:10]:
            print(f"  - {user_id}")
        if len(only_es) > 10:
            print(f"  ... 还有 {len(only_es) - 10} 个")

    if field_diff:
        print(f"\n⚠️  字段不一致的用户 ({len(field_diff)} 个):")
        for user_id, diffs in list(field_diff.items())[:5]:
            print(f"  - {user_id}:")
            for diff in diffs:
                print(f"      {diff}")
        if len(field_diff) > 5:
            print(f"  ... 还有 {len(field_diff) - 5} 个")

    # 修复逻辑
    if args.fix or args.fix_missing or args.fix_extra or args.fix_fields:
        print(f"\n🔧 开始修复不一致数据...")
        mysql_users = checker.fetch_mysql_users()
        total_fixed = 0

        if (args.fix or args.fix_missing) and only_mysql:
            print(f"\n修复 ES 缺失数据...")
            total_fixed += checker.fix_missing_in_es(only_mysql, mysql_users)

        if (args.fix or args.fix_extra) and only_es:
            print(f"\n删除 ES 多余数据...")
            total_fixed += checker.fix_extra_in_es(only_es)

        if (args.fix or args.fix_fields) and field_diff:
            print(f"\n修复字段不一致...")
            total_fixed += checker.fix_field_mismatches(field_diff, mysql_users)

        print(f"\n✅ 修复完成，共修复 {total_fixed} 条记录")
    else:
        if only_mysql or only_es or field_diff:
            print(f"\n💡 提示：使用 --fix 参数可自动修复不一致数据")
            print(f"    或使用 --fix-missing / --fix-extra / --fix-fields 选择性修复")

    checker.close()

    if only_mysql or only_es or field_diff:
        return 1
    else:
        print("\n✅ 数据一致性检查通过！")
        return 0


if __name__ == '__main__':
    sys.exit(main())


